package com.hzxgs.tracelogstarter.trace;

import com.hzxgs.tracelogstarter.properties.TraceProperties;
import com.hzxgs.tracelogstarter.util.TraceUtils;
import io.netty.buffer.ByteBuf;
import io.netty.channel.Channel;
import io.netty.channel.ChannelHandler;
import io.netty.channel.ChannelHandlerContext;
import io.netty.channel.SimpleChannelInboundHandler;
import lombok.extern.slf4j.Slf4j;
import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
import org.springframework.context.annotation.Scope;
import org.springframework.stereotype.Component;

import javax.annotation.Resource;

@Component
@ChannelHandler.Sharable
@ConditionalOnClass(value = { SimpleChannelInboundHandler.class })
@Scope("prototype")
@Slf4j
public class NettyClientTraceServerInHandler extends SimpleChannelInboundHandler {

    @Resource
    private TraceProperties traceProperties;

    @Override
    public void channelInactive(ChannelHandlerContext ctx) {
        if (traceProperties.getEnable()) {
            // 移除channel里的链路追踪信息
            Channel channel = ctx.channel();
            channel.attr(TraceUtils.TRACE_ID_ATTR_KEY).remove();
            channel.attr(TraceUtils.SPAN_ID_ATTR_KEY).remove();
        }
        ctx.pipeline().remove(this);
        ctx.deregister();
        ctx.close();
        ctx.fireChannelInactive();
    }

    @Override
    protected void channelRead0(ChannelHandlerContext ctx, Object msg) throws Exception {
        if (traceProperties.getEnable()) {
            // 从channel里读取发送消息前的traceId
            Channel channel = ctx.channel();
            String traceId = channel.attr(TraceUtils.TRACE_ID_ATTR_KEY).get();
            String spanId = channel.attr(TraceUtils.SPAN_ID_ATTR_KEY).get();
            // 强制更新本次线程的traceId
            TraceUtils.putTrace(traceId, spanId);
        }
        // 不设置时，在返回时会报错
        ByteBuf buf = (ByteBuf) msg;
        buf.retain();
        ctx.fireChannelRead(msg);
    }

    @Override
    public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) {
        if (traceProperties.getEnable()) {
            // 移除channel里的链路追踪信息
            Channel channel = ctx.channel();
            channel.attr(TraceUtils.TRACE_ID_ATTR_KEY).remove();
            channel.attr(TraceUtils.SPAN_ID_ATTR_KEY).remove();
        }

        cause.printStackTrace();
        ctx.pipeline().remove(this);
        ctx.deregister();
        ctx.close();
        ctx.fireExceptionCaught(cause);
    }
}
